import scipy.io.wavfile as wav
import matplotlib.pyplot as plt
from scipy import signal
import numpy as np
import scipy
import wave
import os
import math

fs = 44100
# filePath = 'ae20k.wav'
sonicFreq = 19000
lipOffset = 40
nfft = 8192

# fs, signalWav = wav.read(fileName)
# signalWav = np.fft.fft(signalWav)
# signalWav = np.abs(signalWav)
# time = np.linspace(0, fs, len(signalWav))
# plt.plot(time, signalWav)
# plt.show()

def SaveWav(savePath, signalWav, fs):
    f = wave.open(savePath, "wb")
    # 配置声道数、量化位数和取样频率
    f.setnchannels(1)
    f.setsampwidth(2)
    f.setframerate(fs)
    # 将wav_data转换为二进制数据写入文件
    f.writeframes(signalWav.tostring())
    f.close()

def FiltWav(signalWav, N, lowPass, highPass):
    b, a = signal.butter(N=N, Wn=[lowPass, highPass], btype='bandpass')
    filtSignal = signal.filtfilt(b, a, signalWav)  # data为要过滤的信号
    return filtSignal

def MultiVector(A, B):
    C = np.zeros(len(A))
    for i in range(len(A)):
        C[i] = A[i]*B[i]
    return sum(C)

def LMS(xn, M, mu, signalNum):
    en = np.zeros(signalNum)
    W = [0] * M
    for k in range(signalNum)[M-1:signalNum]:
        x = xn[k-M+1:k+1]
        # d = x.mean()
        y = MultiVector(W,x)
        en[k] = xn[k]-y
        W = np.add(W, mu*en[k]*x)

    yn = np.zeros(len(xn))
    for i in range(M):
        yn[i] = xn[i]
    for k in range(len(xn))[M - 1:len(xn)]:
        x = xn[k - M + 1:k + 1]
        yn[k] = MultiVector(W, x)
    return yn

def SetMax(Zxx, maxThresh):
    shape = Zxx.shape
    for i in range(shape[0]):
        for j in range(shape[1]):
            if Zxx[i][j] > maxThresh:
                Zxx[i][j] = maxThresh
    return Zxx

def Multi2One(Zxx):
    maxData = np.max(abs(Zxx))

    scale = 1.0 / maxData
    Zxx = Zxx * scale
    return Zxx

def DeleteMin(Zxx, minThresh):
    shape = Zxx.shape
    for i in range(shape[0]):
        for j in range(shape[1]):
            if Zxx[i][j] < minThresh:
                Zxx[i][j] = 0
    return Zxx

def SingleSTFT(filePath):
    fs, signalWav = wav.read(filePath)
    cutPoint = fs // 4
    # signalWav = signalWav[cutPoint:len(signalWav)-cutPoint]
    lowPass = 2 * (sonicFreq - lipOffset) / fs
    highPass = 2 * (sonicFreq + lipOffset) / fs

    filtSignal = signalWav
    filtSignal = FiltWav(signalWav=signalWav, N=4, lowPass=lowPass, highPass=highPass)

    filtSignal = filtSignal / max(abs(filtSignal))

    signalNum = len(filtSignal)

    # plt.plot(filtSignal)
    # plt.show()

    # freqAxis = np.linspace(0, fs, signalNum)
    # filtSignal = np.fft.fft(filtSignal)
    # filtSignal = abs(filtSignal)
    # plt.plot(freqAxis, filtSignal)
    # plt.show()

    M = 8
    mu = 0.05

    # lmsSignal = LMS(xn=filtSignal, M=M, mu=mu, signalNum=signalNum)
    lmsSignal = filtSignal
    f, t, Zxx = signal.stft(lmsSignal, fs, nperseg=8192, noverlap=7751, nfft=nfft)
    Zxx = abs(Zxx)

    cutNum = 8
    time = len(t)

    return t, f, Zxx

def subAdjacent(Zxx):
    time = Zxx.shape[1]
    result = np.zeros(Zxx.shape)
    for i in range(1, time):
        result[:, i] = Zxx[:,i] - Zxx[:,i-1]
    return result

def ShowGraph(corpus):

    count = 0
    maxCount = 2

    span = fs / nfft
    basicPos = int(sonicFreq / fs * nfft)
    freqOffset = int(lipOffset // span) + 5
    down = basicPos - freqOffset
    up = basicPos + freqOffset

    for root, dirs, files in os.walk(corpus):
        for file in files:
            filePath = os.path.join(corpus, file)
            data = SingleSTFT(filePath)
            t = data[0]
            f = data[1][down:up]
            Zxx = data[2][down:up, :]
            # Zxx = DeleteMin(Zxx, 0.001)

            # Zxx = subAdjacent(Zxx)
            # Zxx = Normal(Zxx)

            # f = data[1]
            # Zxx = data[2]
            # cmap = plt.get_cmap('Greys')
            cutNum = 8
            time = len(t)
            # t = t[:time-2*cutNum]
            # Zxx = Zxx[:,cutNum:time-cutNum]
            plt.subplot(211+count)
            # plt.title(file.split('_')[0])
            plt.title(file)
            plt.pcolormesh(t, f, Zxx, cmap='YlOrRd')
            count += 1
            if count>=maxCount:
                plt.show()
                count = 0

def ShowVar(corpus):

    count = 0
    maxCount = 2

    span = fs / nfft
    basicPos = int(sonicFreq / fs * nfft)
    freqOffset = int(lipOffset // span) + 5
    down = basicPos - freqOffset
    up = basicPos + freqOffset

    for root, dirs, files in os.walk(corpus):
        for file in files:
            filePath = os.path.join(corpus, file)
            data = SingleSTFT(filePath)
            t = data[0]
            f = data[1][down:up]
            Zxx = data[2][down:up, :]
            # Zxx = DeleteMin(Zxx, 0.001)

            # Zxx = subAdjacent(Zxx)
            # Zxx = Normal(Zxx)

            # f = data[1]
            # Zxx = data[2]
            # cmap = plt.get_cmap('Greys')
            cutNum = 8
            time = len(t)
            t = t[:time-2*cutNum]
            Zxx = Zxx[:,cutNum:time-cutNum]

            var = CalVal(Zxx)
            var = GaosiSmooth(var, 20, 1)
            var = ThreshVar(var, 3*10e-8)

            plt.subplot(211+count)
            # plt.title(file.split('_')[0])
            plt.title(file)
            plt.plot(t[:], var)
            count += 1
            if count>=maxCount:
                plt.show()
                count = 0


def SignalGrad(Zxx):
    time = Zxx.shape[1]
    dealed = []
    for i in range(1, time):
        temp = Zxx[:, i] - Zxx[:, i-1]
        dealed.append(temp)
    dealed = np.array(dealed)
    dealed = np.moveaxis(dealed, 0, 1)
    return dealed

def TreshSignal(Zxx, thresh):
    shape = Zxx.shape
    for i in range(shape[1]):
        for j in range(shape[0]):
            if Zxx[j][i] < thresh:
                Zxx[j][i] = 0
    return Zxx

def ThreshVar(var, thresh):
    for i in range(len(var)):
        if var[i] <= thresh:
            var[i] = 0
    return var

def Normal(Zxx):
    mean = np.mean(Zxx, axis=1)
    std = np.std(Zxx, axis=1)
    yShape = Zxx.shape[1]
    for i in range(yShape):
        Zxx[:,i] = (Zxx[:,i]-mean)/std
    Zxx[np.isnan(Zxx)] = 0
    return Zxx

def CalVal(Zxx):
    var = np.var(Zxx, axis=0)
    # index = range(0, Zxx.shape[1])
    # dian = []
    # for temp in zip(index, var):
    #     dian.append(temp)
    # dian = np.array(dian)
    return var[:]

def GaussWeight(x, sigma):
    w = 1 / ((2 * np.pi) ** 0.5) * np.exp(-x*x/2/sigma/sigma)
    return w

def GaussSmooth(var, lenNum, sigma):
    varNum = len(var)
    varBefore = var[lenNum-1::-1]
    dataAfter = var[sigma-lenNum:]
    varAfter = dataAfter[::-1]
    var = np.concatenate((varBefore, var, varAfter))
    mid = lenNum//2
    coreIndex = [i - mid for i in range(lenNum)]
    coreIndex = np.array(coreIndex)
    w = GaussWeight(coreIndex, sigma)
    result = []
    # for i in range(mid, varNum-mid):
    #     temp = w*var[i-mid, i+]

    for i in range(lenNum, varNum+lenNum):
        temp = np.dot(w,var[i-mid:i+lenNum-mid])
        result.append(temp)
    return np.array(result)

def SplitWordIndex(var, noiseThresh, peakThresh, countThresh, winLen):
    data = []
    count = 0
    lastIndex = 0
    for i in range(len(var)-winLen):
        temp = var[i:i+winLen]
        temp = (temp > noiseThresh)
        # print(temp)
        if True not in temp:
            count += 1
        else:
            count = 0

        if count >= countThresh:
            # print(count)
            active = var[lastIndex:i]
            peak = signal.find_peaks(active)
            if len(peak[0]) >0:
                judge = (peak[0] > peakThresh)
            else:
                judge = np.array([False])

            if True in judge:

                data.append((lastIndex, i))
                lastIndex = i
            count = 0

    return data

def SplitWord(Zxx, var, noiseThresh, peakThresh, countThresh, winLen):
    data = []
    indexData = SplitWordIndex(var, noiseThresh, peakThresh, countThresh, winLen)
    print(indexData)
    for index in indexData:
        # temp = var[index[0]:index[1]]
        temp = Zxx[:, index[0]:index[1]]
        data.append(temp)
    return data

def ShowCompare(corpus):
    pathList = []
    basePath = None
    fs = 44100
    lowPass = 2 * (sonicFreq - lipOffset) / fs
    highPass = 2 * (sonicFreq + lipOffset) / fs

    for root, dirs, files in os.walk(corpus):
        for file in files:
            filePath = os.path.join(corpus, file)
            if file == 'idle.wav':
                basePath = filePath
                continue
            pathList.append(filePath)
    for i in range(len(pathList)):
        N = 5
        fs, signal1 = wav.read(basePath)
        fs, signal2 = wav.read(pathList[i])
        len1 = len(signal1)
        len2 = len(signal2)
        time1 = np.linspace(0, len1/fs, len1)
        time2 = np.linspace(0, len2/fs, len2)
        # signal1 = FiltWav(signal1, N, lowPass, highPass)
        # signal2 = FiltWav(signal2, N, lowPass, highPass)
        plt.subplot(211)
        plt.title(basePath.split(os.sep)[-1])
        plt.plot(time1, signal1)
        plt.subplot(212)
        plt.title(pathList[i].split(os.sep)[-1])
        plt.plot(time2, signal2)
        plt.show()


corpus = 'haytest'
ShowGraph(corpus)
ShowVar(corpus)

# corpus = 'hand'
# ShowCompare(corpus)


f1 = 'haytest/hay_8.wav'
fs = 44100
# filePath = 'ae20k.wav'
lipOffset = 40
nfft = nfft
lowPass = 2 * (sonicFreq - lipOffset) / fs
highPass = 2 * (sonicFreq + lipOffset) / fs



# fs, signal1 = wav.read(f1)
# fs, signal2 = wav.read(f2)
#
# signal1 = FiltWav(signal1, 3, lowPass, highPass)
# signal2 = FiltWav(signal2, 3, lowPass, highPass)
#
# plt.subplot(211)
# plt.plot(signal1)
# plt.subplot(212)
# plt.plot(signal2)
# plt.show()


span = fs/nfft
basicPos = int(sonicFreq/fs*nfft)
freqOffset = int(lipOffset//span) + 10
down = basicPos - freqOffset
up = basicPos + freqOffset

t, f, Zxx = SingleSTFT(f1)

timeNum = len(t)
cutPoint = 8

t = t[cutPoint:timeNum-cutPoint]
f = f[down:up]
Zxx = Zxx[down:up, cutPoint:timeNum-cutPoint-2]

# Zxx = DeleteMin(Zxx, 0.01)
# Zxx = Normal(Zxx)
Zxx = subAdjacent(Zxx)
# Zxx = Multi2One(Zxx)

vars = CalVal(Zxx)
vars = GaosiSmooth(vars, 20, 4)
plt.plot(vars)
plt.show()

noiseThresh = 3*10e-8
peakThresh = 7*10e-8
countThresh = 3
winLen = 15
wordData = SplitWord(Zxx=Zxx, var=vars, noiseThresh=noiseThresh,
                     peakThresh=peakThresh, countThresh=countThresh, winLen=winLen)

for Zxx in wordData:
    tLen = Zxx.shape[1]
    t = np.linspace(0, 0.1*tLen, tLen)
    plt.pcolormesh(t, f, Zxx, cmap='YlOrRd')
    plt.show()
    # plt.plot(Zxx)
    # plt.show()



# plt.pcolormesh(t, f, Zxx, cmap='YlOrRd')
# Zxx = TreshSignal(Zxx, 0.08)
# plt.pcolormesh(t, f, Zxx, cmap='YlOrRd')
# plt.pcolormesh(t, f, Zxx, cmap='YlOrRd')
# plt.title('STFT Magnitude')
# plt.ylabel('Frequency [Hz]')
# plt.xlabel('Time [sec]')
# plt.show()


def takeSecond(elem):
    return elem[1]

def GetPeak(curFreq):
    peakArr = []
    curLen = len(curFreq)
    for i in range(1, curLen-1):
        if curFreq[i]>curFreq[i-1] and curFreq[i]>curFreq[i+1]:
            peakArr.append((i, curFreq[i]))
    if len(peakArr) < 2:
        peakArr.append((0,0))
        peakArr.append((0,0))
    return peakArr

# Zxx = np.abs(Zxx)

# Zxx = SetMax(Zxx, 0.05)
freqVary = []
thresh = 0.12
threshSecond = 0.11
centerIndex = []

# f = f[down:up]

# for i in range(Zxx.shape[1]):
#     freq = Zxx[:,i]
#     # print(freq)
#     plt.plot(f, freq)
#     plt.show()


